import numpy as np
from fastapi.staticfiles import StaticFiles
from fastapi import FastAPI, Request
from fastapi.responses import HTMLResponse, StreamingResponse
import os
import sys
import asyncio
import pyaudio
import time
import dashscope
import multiprocessing

from starlette.websockets import WebSocketDisconnect
try:
    import sounddevice as sd
except ImportError as e:
    print("Please install sounddevice first. You can use")
    print()
    print("  pip install sounddevice")
    print()
    print("to install it")
    sys.exit(-1)
from get_config import config  # 从 get_config 中加载配置
import json
import logging
import threading
from fastapi.websockets import WebSocket
from concurrent.futures import ThreadPoolExecutor
from test import AudioRecognition


# FastAPI 初始化
app = FastAPI()

# 配置路径
BASE_DIR = os.path.dirname(os.path.abspath(__file__))
app.mount("/static", StaticFiles(directory=os.path.join(BASE_DIR, "static")), name="static")

# 从配置文件加载音频和识别参数
audio_config = config["audio"]
recognition_config = config["recognition"]

# dashscope.api_key=recognition_config['api_key']
print("启动服务中")
# PyAudio 初始化
audio = pyaudio.PyAudio()

# 自动获取硬件设备参数
default_device_info = audio.get_default_input_device_info()
RATE = int(audio_config["rate"])  # 硬件采样率
CHANNELS = default_device_info["maxInputChannels"]  # 通道数
FORMAT = pyaudio.paInt16  # 默认设置为 16 位深度
TARGET_RATE = audio_config["target_rate"]  # 配置中的目标采样率
CHUNK = audio_config["chunk"]  # 每帧大小

# STT 配置
STT_MODEL = recognition_config["model"]

# 打印当前设置
print("音频设备信息:")
print(f"设备名称: {default_device_info['name']}")
print(f"硬件采样率: {RATE}")
print(f"输入通道数: {CHANNELS}")
print(f"目标采样率: {TARGET_RATE}")
print(f"帧大小: {CHUNK}")
print("STT 配置:")
print(f"识别模型: {STT_MODEL}")

def remove_low_db(samples, db_threshold):
    """
    将低于指定分贝阈值的音频样本置零。

    参数：
    - samples: 音频样本，numpy 数组
    - db_threshold: 分贝阈值，低于该值的样本将被置零
    """
    safe_samples = np.where(samples == 0, 1e-10, samples)
    db_values = 20 * np.log10(np.abs(safe_samples))
    samples[db_values < db_threshold] = 0
    return samples


@app.get("/test")
async def get():
    """
    测试静态页面
    """
    file_path = os.path.join(BASE_DIR, "static", "index.html")

    # 验证文件是否存在
    if not os.path.exists(file_path):
        return HTMLResponse("<h1>Error: index.html not found</h1>", status_code=404)

    # 返回 HTML 页面
    with open(file_path, "r", encoding="utf-8") as f:
        html_content = f.read()
    return HTMLResponse(html_content)


# audio_recognition = AudioRecognition(
#         model=STT_MODEL,
#         format=FORMAT_STT,
#         sample_rate=TARGET_RATE,
#         chunk_size=CHUNK,
# )
# executor = ThreadPoolExecutor()

# @app.on_event("startup")
async def startup_event():
    """
    在异步线程中启动 AudioRecognition 的识别任务。
    """
    loop = asyncio.get_event_loop()

    # 定义阻塞任务的包装器
    def run_audio_recognition():
        # asyncio.run(audio_recognition.start_recognition())
        logging.info("")

    # 提交任务到线程池，并确保它不会阻塞主线程
    #loop.run_in_executor(executor, run_audio_recognition)
    logging.info("音频识别任务已提交到线程池运行。")
# @app.on_event("shutdown")
async def shutdown_event():
    """
    在服务器停止时释放资源并终止音频识别任务。
    """
    logging.info("服务器停止中，释放资源...")
    try:
        # wait audio_recognition.stop_recognition()  # 停止音频识别任务
        logging.info("音频识别任务已停止，资源已释放。")
    except Exception as e:
        logging.error(f"在释放资源时发生错误: {e}")
# @app.websocket("/ws")
async def websocket_endpoint(websocket: WebSocket):
    await websocket.accept()
    logging.info("WebSocket连接已建立")
    await startup_event()
    stop_event = threading.Event()
    def consumer(queue, websocket, stop_event):
        """
        消费队列中的数据并发送到 WebSocket。
        """
        while not stop_event.is_set():  # 检查停止标志
            if not queue.empty():
                data = queue.get()
                try:
                    if isinstance(data, dict):
                        message = json.dumps(data)
                    else:
                        message = str(data)
                    asyncio.run(websocket.send_text(message))
                except Exception as e:
                    logging.error(f"发送到 WebSocket 时出错: {e}")
                    break
            else:
                time.sleep(0.1)  # 队列为空时稍作休眠

    # 启动消费线程
    #consumer_thread = threading.Thread(target=consumer, args=(audio_recognition.queue, websocket,stop_event))
    #consumer_thread.start()

    try:
        # 保持 WebSocket 打开状态
        while True:
            await websocket.receive_text()
    except WebSocketDisconnect:
        logging.info("WebSocket连接已断开")
    except Exception as e:
        logging.error(f"WebSocket发生错误: {e}")
    finally:
        # 停止消费线程
        logging.info("释放资源...")
        stop_event.set()  # 设置停止标志
        #consumer_thread.join()  # 等待线程结束
        await shutdown_event()
        # audio_recognition.stop_recognition()  # 停止音频识别任务
        logging.info("资源释放完成")

from tts import recognizer


async def send_message(websocket, msg):
    """异步发送 WebSocket 消息"""
    await websocket.send_text(json.dumps(msg))


@app.websocket("/ws2")
async def websocket_endpoint2(websocket: WebSocket):
    await websocket.accept()
    stop_event = threading.Event()

    loop = asyncio.get_event_loop()
    def consumer(websocket, stop_event):
        recognizer.reset()  # 重置识别器状态
        with sd.InputStream(channels=1, dtype="float32", samplerate=RATE) as s:
            while not stop_event.is_set():
                try:
                    # 从麦克风读取音频数据
                    samples, _ = s.read(CHUNK)  # a blocking read
                    samples = samples.reshape(-1)
                    # 将音频数据传递给识别器
                    recognizer.accept_waveform(RATE, samples)
                    # 进行解码
                    result = recognizer.text
                    if result and not recognizer.is_endpoint:
                        msg = {"event": "result", "text": result}
                        logging.info(msg)
                        # asyncio.run(websocket.send_text(json.dumps(msg)))
                        loop.create_task(websocket.send_text(json.dumps(msg)))
                    if recognizer.is_endpoint and result:
                        msg = {"event": "endpoint", "text": result}
                        # asyncio.run(websocket.send_text(json.dumps(msg)))
                        loop.create_task(websocket.send_text(json.dumps(msg)))
                        logging.info(msg)
                        recognizer.reset()
                except Exception as e:
                    logging.error(f"发送到 WebSocket 时出错: {e}")
                    break
            s.close()

    consumer_thread = threading.Thread(target=consumer, args=(websocket, stop_event))
    consumer_thread.start()
    try:
        # 保持 WebSocket 打开状态
        while True:
            await websocket.receive_text()
    except WebSocketDisconnect:
        logging.info("WebSocket连接已断开")
    except Exception as e:
        logging.error(f"WebSocket发生错误: {e}")
    finally:
        # 停止消费线程
        logging.info("释放资源...")
        stop_event.set()  # 设置停止标志
        consumer_thread.join()  # 等待线程结束
        try:
            await websocket.close()
        except:
            pass
        logging.info("资源释放完成")

def consumer_process(websocket_queue, stop_event):
    """消费者进程，处理音频数据并将消息发送到主进程"""
    recognizer.reset()  # 重置识别器状态
    with sd.InputStream(channels=1, dtype="float32", samplerate=RATE) as s:
        while not stop_event.is_set():
            try:
                # 从麦克风读取音频数据
                samples, _ = s.read(CHUNK)  # a blocking read
                samples = samples.reshape(-1)
                # 将音频数据传递给识别器
                recognizer.accept_waveform(RATE, samples)
                # 进行解码
                result = recognizer.text
                if result and not recognizer.is_endpoint:
                    msg = {"event": "result", "text": result}
                    logging.info(msg)
                    websocket_queue.put(msg)  # 将消息放入队列中，主进程会发送
                if recognizer.is_endpoint and result:
                    msg = {"event": "endpoint", "text": result}
                    logging.info(msg)
                    websocket_queue.put(msg)  # 将消息放入队列中，主进程会发送
                    recognizer.reset()
            except Exception as e:
                logging.error(f"发送到 WebSocket 时出错: {e}")
                break
        s.close()

from queue import Empty
def consumer_task(websocket_queue, stop_event, websocket,loop):
    """线程任务，异步从队列中获取数据并发送"""
    while not stop_event.is_set():
        try:
            # 异步从队列中获取消息
            msg = websocket_queue.get()  # 阻塞等待队列中的消息
            if msg:
                loop.create_task(send_message(websocket, msg))  # 异步发送消息
        except Empty:
            pass  # 队列为空，继续等待
        except Exception as e:
            logging.error(f"消费者线程发生错误: {e}")
            break
@app.websocket("/ws3")
async def websocket_endpoint3(websocket: WebSocket):
    await websocket.accept()
    stop_event = multiprocessing.Event()
    websocket_queue = multiprocessing.Queue()
    loop = asyncio.get_event_loop()  # 获取事件循环

    # 启动生产者进程
    process = multiprocessing.Process(target=consumer_process, args=(websocket_queue, stop_event))
    process.start()

    consumer_thread = threading.Thread(target=consumer_task, args=(websocket_queue, stop_event, websocket,loop))
    consumer_thread.start()

    try:
        # 保持 WebSocket 打开状态
        while True:
            await websocket.receive_text()
    except WebSocketDisconnect:
        logging.info("WebSocket连接已断开")
    except Exception as e:
        logging.error(f"WebSocket发生错误: {e}")
    finally:
        # 停止消费者进程
        logging.info("释放资源...")
        stop_event.set()  # 设置停止标志
        process.join()  # 等待进程结束
        try:
            await websocket.close()
        except:
            pass
        logging.info("资源释放完成")


